#include <netdb.h>
#include <netinet/in.h>
#include <sys/socket.h>

#include <iostream>
#include <cstring>
#include <cerrno>
#include <cassert>

#include <cstdio>

#include "utils.h"
#include "sgconf.h"
#include "proto.h"

void DumpCrypto(const char * data, size_t size)
{
    std::string dumpstr;
    for (unsigned i = 0; i < size; ++i) {
        char buf[256];
        snprintf(buf, 256, "%x", *(data + i));
        dumpstr += buf;
    }

    std::cerr << "Crypto dump: '" << dumpstr << "'" << std::endl;
}

SGCONF::SGCONF()
    : sd(-1),
      hostOk(false)
{
    iv = new unsigned char[8];
    memset(iv, 0, 8);
}

SGCONF::~SGCONF()
{
    delete[] iv;
    close(sd);
}

bool SGCONF::SetHost(const std::string & h)
{
    struct addrinfo * ai;

    int res = getaddrinfo(h.c_str(), NULL, NULL, &ai);

    if (res != 0) {   
        std::cerr << "error in getaddrinfo: " << gai_strerror(res) << "\n";
        return false;
    }

    memcpy(&remoteAddr, ai->ai_addr, sizeof(remoteAddr));

    freeaddrinfo(ai);

    std::cout << "Host ip: " << inet_ntostring(remoteAddr.sin_addr.s_addr) << std::endl;

    hostOk = true;

    return true;
}

bool SGCONF::Perform(const std::string & xml)
{
    if (!Connect()) {
        return false;
    }

    struct REQ::HEADER reqHdr;

    strncpy(reqHdr.magic, PROTO_MAGIC, sizeof(reqHdr.magic));
    reqHdr.version = htonl(2 << 8 | 0);
    memset(reqHdr.login, 0, sizeof(reqHdr.login));
    strncpy(reqHdr.login, login.c_str(), sizeof(reqHdr.login));

    struct REQ::CRYPTO_HEADER reqCryptoHdr;
    memset(reqCryptoHdr.login, 0, sizeof(reqCryptoHdr.login));
    strncpy(reqCryptoHdr.login, login.c_str(), sizeof(reqCryptoHdr.login));
    if (xml.size() % 8 == 0) {
        reqCryptoHdr.dataSize = htonl(xml.size());
    } else {
        reqCryptoHdr.dataSize = htonl((xml.size() / 8 + 1) * 8);
    }

    size_t size = sizeof(reqHdr);
    if (!WriteBlock(reinterpret_cast<char *>(&reqHdr), size, 5000)) {
        return false;
    }

    BF_set_key(&key, password.length(), reinterpret_cast<const unsigned char *>(password.c_str()));

    unsigned char buffer[sizeof(reqCryptoHdr)];
    size = sizeof(reqCryptoHdr);

    BF_cbc_encrypt(reinterpret_cast<const unsigned char *>(&reqCryptoHdr), buffer, size, &key, iv, BF_ENCRYPT);

    if (!WriteBlock(buffer, size, 5000)) {
        return false;
    }

    return WriteXML(xml);
}

bool SGCONF::ReadBlock(void * dest, size_t & size, int timeout) const
{
    char * ptr = static_cast<char *>(dest);
    unsigned readSize = 0;
    while (readSize < size) {
        struct timeval tv;
        tv.tv_sec = 0;
        tv.tv_usec = timeout * 1000;

        fd_set rfds;
        FD_ZERO(&rfds);
        FD_SET(sd, &rfds);

        int res = select(sd + 1, &rfds, NULL, NULL, &tv);
        /* Don't rely on the value of tv now! */

        if (res < 0) {
            std::cerr << "SGCONF::ReadBlock() Select error: '" << strerror(errno) << "'\n";
            return false;
        }

        if (res == 0) {
            // Timeout
            size = readSize;
            return false;
        }

        res = read(sd, ptr + readSize, size - readSize);

        if (res == 0) {
            std::cerr << "SGCONF::ReadBlock() EOF\n";
            return false;
        }

        // Ignore 'Interrupted system call' errors
        if (res < 0) {
            if (errno != EINTR) {
                std::cerr << "SGCONF::ReadBlock() Read error: '" << strerror(errno) << "'\n";
                return false;
            }
            continue;
        }

        readSize += res;
    }

    return true;
}

bool SGCONF::WriteBlock(const void * source, size_t & size, int timeout) const
{
    const char * ptr = static_cast<const char *>(source);
    unsigned writeSize = 0;
    while (writeSize < size) {
        struct timeval tv;
        tv.tv_sec = 0;
        tv.tv_usec = timeout * 1000;

        fd_set wfds;
        FD_ZERO(&wfds);
        FD_SET(sd, &wfds);

        int res = select(sd + 1, NULL, &wfds, NULL, &tv);
        /* Don't rely on the value of tv now! */

        if (res < 0) {
            std::cerr << "SGCONF::WriteBlock() Select error: '" << strerror(errno) << "'\n";
            return false;
        }

        if (res == 0) {
            // Timeout
            size = writeSize;
            return false;
        }

        res = write(sd, ptr + writeSize, size - writeSize);

        // Ignore 'Interrupted system call' errors
        if (res < 0) {
            if (errno != EINTR) {
                std::cerr << "SGCONF::WriteBlock() Write error: '" << strerror(errno) << "'\n";
                return false;
            }
            continue;
        }

        writeSize += res;
    }

    return true;
}

bool SGCONF::Connect()
{
    if (!hostOk) {
        return false;
    }

    sd = socket(AF_INET, SOCK_STREAM, 0);

    if (sd < 0) {
        std::cerr << "SGCONF::Connect() Socket error: " << strerror(errno) << std::endl;
        return false;
    }

    remoteAddr.sin_family = AF_INET;

    if (connect(sd, (struct sockaddr *)&remoteAddr, sizeof(remoteAddr))) {
        std::cerr << "SGCONF::Connect() Connect error: " << strerror(errno) << std::endl;
        return false;
    }

    return true;
}

bool SGCONF::WriteXML(const std::string & xml) const
{
    size_t pos = 0;
    std::string data(xml);
    if (data.size() % 8) {
        size_t delta = (data.size() / 8 + 1) * 8 - data.size();
        data.append(delta, ' ');
    }
    while (pos < data.size()) {
        unsigned char source[1496];
        unsigned char buffer[1496];

        size_t size;
        if (data.size() - pos > sizeof(source)) {
            memcpy(source, data.c_str() + pos, sizeof(source));
            size = sizeof(source);
        } else {
            memset(source, 0, sizeof(source));
            memcpy(source, data.c_str() + pos, data.size() - pos);
            size = data.size() - pos;
        }

        BF_cbc_encrypt(source, buffer, size, &key, iv, BF_ENCRYPT);

        if (!WriteBlock(buffer, size, 5000)) {
            return false;
        }

        pos += size; // size?
    }

    return true;
}

bool SGCONF::GetAnswer(std::string & xml) const
{
    struct RESP::HEADER respHeader;

    size_t size = sizeof(respHeader);
    if (!ReadBlock(&respHeader, size, 5000)) {
        std::cerr << "No answer header within 5 sec\n";
        return false;
    }

    if (strncmp(respHeader.magic, PROTO_MAGIC, sizeof(respHeader.magic))) {
        std::cerr << "Invalid magic code in header\n";
        return false;
    }

    uint32_t version = ntohl(respHeader.version);
    if (version > (2 << 8 | 0)) {
        std::cerr << "Unsupported version\n";
        return false;
    }

    // Check code

    unsigned char buffer[sizeof(struct RESP::CRYPTO_HEADER)];
    size = sizeof(struct RESP::CRYPTO_HEADER);

    if (!ReadBlock(buffer, size, 5000)) {
        std::cerr << "No crypto header within 5 secs\n";
        return false;
    }

    struct REQ::CRYPTO_HEADER respCryptoHeader;

    BF_cbc_encrypt(buffer, reinterpret_cast<unsigned char *>(&respCryptoHeader), sizeof(struct REQ::CRYPTO_HEADER), &key, iv, BF_DECRYPT);

    respCryptoHeader.login[sizeof(respCryptoHeader.login) - 1] = 0;

    std::string cryptoLogin(respCryptoHeader.login);

    if (login != cryptoLogin) {
        std::cerr << "Password is invalid\n";
        return false;
    }

    assert(respCryptoHeader.dataSize % 8 == 0);

    xml = "";

    char block[1496];
    unsigned char cryptoBlock[1496];
    size_t length = 0;
    uint32_t dataSize = ntohl(respCryptoHeader.dataSize);

    while (length < dataSize) {
        size_t delta = dataSize - length;
        if (delta > sizeof(cryptoBlock)) {
            delta = sizeof(cryptoBlock);
        }
        size_t bs = delta;
        ReadBlock(cryptoBlock, bs, 5000);
        if (bs != delta) {
            std::cerr << "No data within 5 secs\n";
            return false;
        }

        BF_cbc_encrypt(cryptoBlock, reinterpret_cast<unsigned char *>(block), bs, &key, iv, BF_DECRYPT);

        xml.append(block, bs);

        length += bs;
    }

    return true;
}
